"""
连接操作符实现
"""

from typing import Optional, List, Dict, Any, Iterator, Set, Tuple
import time
import threading
from collections import defaultdict

from ..executor_base import IteratorExecutor
from ..types import (
    Record, RecordBatch, ExecutionContext, JoinType, DataValue
)

class JoinOperatorBase(IteratorExecutor):
    """连接操作符基类"""
    
    def __init__(self, executor_id: str, context: ExecutionContext,
                 left_child: IteratorExecutor, right_child: IteratorExecutor,
                 join_type: JoinType, join_condition: str):
        super().__init__(executor_id, context)
        self.left_child = left_child
        self.right_child = right_child
        self.join_type = join_type
        self.join_condition = join_condition
        
        # 解析连接条件
        self.left_keys, self.right_keys = self._parse_join_condition(join_condition)
        
        # 添加子执行器
        self.add_child(left_child)
        self.add_child(right_child)
    
    def _parse_join_condition(self, condition: str) -> Tuple[List[str], List[str]]:
        """解析连接条件"""
        # 简化的条件解析，实际应该支持更复杂的表达式
        # 假设条件格式为 "left.col1 = right.col2 AND left.col3 = right.col4"
        left_keys = []
        right_keys = []
        
        if "=" in condition:
            parts = condition.split(" AND ")
            for part in parts:
                if "=" in part:
                    left_expr, right_expr = part.split("=", 1)
                    left_expr = left_expr.strip()
                    right_expr = right_expr.strip()
                    
                    # 提取列名（去掉表前缀）
                    if "." in left_expr:
                        left_col = left_expr.split(".", 1)[1]
                    else:
                        left_col = left_expr
                    
                    if "." in right_expr:
                        right_col = right_expr.split(".", 1)[1]
                    else:
                        right_col = right_expr
                    
                    left_keys.append(left_col)
                    right_keys.append(right_col)
        
        return left_keys, right_keys
    
    def _extract_join_key(self, record: Record, key_columns: List[str]) -> Tuple:
        """提取连接键"""
        key_values = []
        for col in key_columns:
            key_values.append(record.get(col))
        return tuple(key_values)
    
    def _merge_records(self, left_record: Record, right_record: Record) -> Record:
        """合并两条记录"""
        merged = {}
        
        # 获取表别名，优先使用table_alias，然后是table_name
        left_table = getattr(self.left_child, 'table_alias', 
                            getattr(self.left_child, 'table_name', 'left'))
        right_table = getattr(self.right_child, 'table_alias', 
                             getattr(self.right_child, 'table_name', 'right'))
        
        # 添加左表记录，直接使用表名作为前缀
        for key, value in left_record.items():
            merged[f"{left_table}.{key}"] = value
        
        # 添加右表记录，直接使用表名作为前缀
        for key, value in right_record.items():
            merged[f"{right_table}.{key}"] = value
        
        return merged
    
    def _create_null_record(self, template_record: Record, prefix: str) -> Record:
        """创建空记录"""
        null_record = {}
        for key in template_record.keys():
            null_record[f"{prefix}.{key}"] = None
        return null_record

class HashJoinOperator(JoinOperatorBase):
    """哈希连接操作符"""
    
    def __init__(self, executor_id: str, context: ExecutionContext,
                 left_child: IteratorExecutor, right_child: IteratorExecutor,
                 join_type: JoinType, join_condition: str):
        super().__init__(executor_id, context, left_child, right_child, join_type, join_condition)
        self.hash_table: Dict[Tuple, List[Record]] = defaultdict(list)
        self.build_phase_complete = False
        self.left_records_seen: Set[Tuple] = set()
        self.right_template_record = {}
    
    def open(self):
        """打开哈希连接操作符"""
        with self.lock:
            if self.is_open:
                return
            
            self.is_open = True
            
            # 打开子执行器
            self.left_child.open()
            self.right_child.open()
            
            # 构建哈希表（使用右表）
            self._build_hash_table()
            
            # 创建探测迭代器
            self.current_iterator = self._create_probe_iterator()
            
            self.update_stats(memory_used=self._estimate_hash_table_size())
    
    def _build_hash_table(self):
        """构建哈希表"""
        build_start_time = time.time()
        
        # 读取右表所有记录构建哈希表
        while True:
            right_batch = self.right_child.next_batch()
            if not right_batch:
                break
            
            for record in right_batch:
                # 提取连接键
                join_key = self._extract_join_key(record, self.right_keys)
                self.hash_table[join_key].append(record)
                
                # 保存模板记录用于外连接
                if not self.right_template_record:
                    self.right_template_record = record
            
            self.update_stats(rows_processed=len(right_batch))
        
        self.build_phase_complete = True
        build_time = (time.time() - build_start_time) * 1000
        self.stats.cpu_time_ms += build_time
    
    def _create_probe_iterator(self) -> Iterator[Record]:
        """创建探测迭代器"""
        # 探测阶段：遍历左表
        while True:
            left_batch = self.left_child.next_batch()
            if not left_batch:
                break
            
            for left_record in left_batch:
                self.check_timeout()
                
                # 提取连接键
                left_key = self._extract_join_key(left_record, self.left_keys)
                self.left_records_seen.add(left_key)
                
                # 在哈希表中查找匹配
                matching_records = self.hash_table.get(left_key, [])
                
                if matching_records:
                    # 找到匹配记录
                    for right_record in matching_records:
                        merged = self._merge_records(left_record, right_record)
                        self.update_stats(rows_returned=1)
                        yield merged
                
                elif self.join_type in [JoinType.LEFT, JoinType.FULL]:
                    # 左外连接或全外连接：左表记录没有匹配时仍然输出
                    null_right = self._create_null_record(self.right_template_record, "right")
                    merged = self._merge_records(left_record, {})
                    merged.update(null_right)
                    self.update_stats(rows_returned=1)
                    yield merged
        
        # 右外连接或全外连接：输出右表中没有匹配的记录
        if self.join_type in [JoinType.RIGHT, JoinType.FULL]:
            for right_key, right_records in self.hash_table.items():
                if right_key not in self.left_records_seen:
                    for right_record in right_records:
                        # 创建左表的空记录（如果有左表模板的话）
                        left_template = {}
                        if hasattr(self.left_child, 'table_schema') and self.left_child.table_schema:
                            # 使用左表schema创建空记录
                            for col in self.left_child.table_schema:
                                left_template[col] = None
                        
                        merged = self._merge_records(left_template, right_record)
                        self.update_stats(rows_returned=1)
                        yield merged
    
    def _estimate_hash_table_size(self) -> int:
        """估算哈希表大小"""
        # 简单估算：每个记录约100字节
        total_records = sum(len(records) for records in self.hash_table.values())
        return total_records * 100
    
    def close(self):
        """关闭哈希连接操作符"""
        with self.lock:
            if not self.is_open:
                return
            
            # 关闭子执行器
            self.left_child.close()
            self.right_child.close()
            
            # 清理哈希表
            self.hash_table.clear()
            self.left_records_seen.clear()
            
            self.is_open = False
            self.current_iterator = None

class NestedLoopJoinOperator(JoinOperatorBase):
    """嵌套循环连接操作符"""
    
    def __init__(self, executor_id: str, context: ExecutionContext,
                 left_child: IteratorExecutor, right_child: IteratorExecutor,
                 join_type: JoinType, join_condition: str):
        super().__init__(executor_id, context, left_child, right_child, join_type, join_condition)
        self.right_records: List[Record] = []
        self.current_left_record: Optional[Record] = None
        self.right_index = 0
    
    def open(self):
        """打开嵌套循环连接操作符"""
        with self.lock:
            if self.is_open:
                return
            
            self.is_open = True
            
            # 打开子执行器
            self.left_child.open()
            self.right_child.open()
            
            # 预加载右表数据
            self._preload_right_table()
            
            # 创建连接迭代器
            self.current_iterator = self._create_join_iterator()
    
    def _preload_right_table(self):
        """预加载右表数据"""
        while True:
            right_batch = self.right_child.next_batch()
            if not right_batch:
                break
            
            self.right_records.extend(right_batch)
            self.update_stats(rows_processed=len(right_batch))
    
    def _create_join_iterator(self) -> Iterator[Record]:
        """创建连接迭代器"""
        # 外循环：遍历左表
        while True:
            left_batch = self.left_child.next_batch()
            if not left_batch:
                break
            
            for left_record in left_batch:
                self.check_timeout()
                
                has_match = False
                
                # 内循环：遍历右表
                for right_record in self.right_records:
                    if self._records_match(left_record, right_record):
                        has_match = True
                        merged = self._merge_records(left_record, right_record)
                        self.update_stats(rows_returned=1)
                        yield merged
                
                # 处理外连接情况
                if not has_match and self.join_type in [JoinType.LEFT, JoinType.FULL]:
                    null_right = self._create_null_record(
                        self.right_records[0] if self.right_records else {}, "right"
                    )
                    merged = self._merge_records(left_record, {})
                    merged.update(null_right)
                    self.update_stats(rows_returned=1)
                    yield merged
        
        # 右外连接：输出右表中没有匹配的记录
        if self.join_type in [JoinType.RIGHT, JoinType.FULL]:
            # 重新打开左表以进行匹配检查
            self.left_child.close()
            self.left_child.open()
            
            # 收集所有左表记录用于匹配检查
            all_left_records = []
            while True:
                left_batch = self.left_child.next_batch()
                if not left_batch:
                    break
                all_left_records.extend(left_batch)
            
            for right_record in self.right_records:
                has_match = False
                
                # 检查是否有左表记录匹配
                for left_record in all_left_records:
                    if self._records_match(left_record, right_record):
                        has_match = True
                        break
                
                if not has_match:
                    # 创建左表的空记录
                    left_template = {}
                    if all_left_records:
                        # 使用左表第一条记录作为模板创建空记录
                        for key in all_left_records[0].keys():
                            left_template[key] = None
                    
                    merged = self._merge_records(left_template, right_record)
                    self.update_stats(rows_returned=1)
                    yield merged
    
    def _records_match(self, left_record: Record, right_record: Record) -> bool:
        """检查两条记录是否匹配"""
        for left_key, right_key in zip(self.left_keys, self.right_keys):
            left_value = left_record.get(left_key)
            right_value = right_record.get(right_key)
            
            if left_value != right_value:
                return False
        
        return True
    
    def close(self):
        """关闭嵌套循环连接操作符"""
        with self.lock:
            if not self.is_open:
                return
            
            # 关闭子执行器
            self.left_child.close()
            self.right_child.close()
            
            # 清理数据
            self.right_records.clear()
            
            self.is_open = False
            self.current_iterator = None

class SortMergeJoinOperator(JoinOperatorBase):
    """排序合并连接操作符"""
    
    def __init__(self, executor_id: str, context: ExecutionContext,
                 left_child: IteratorExecutor, right_child: IteratorExecutor,
                 join_type: JoinType, join_condition: str):
        super().__init__(executor_id, context, left_child, right_child, join_type, join_condition)
        self.left_sorted_records: List[Record] = []
        self.right_sorted_records: List[Record] = []
        self.left_index = 0
        self.right_index = 0
    
    def open(self):
        """打开排序合并连接操作符"""
        with self.lock:
            if self.is_open:
                return
            
            self.is_open = True
            
            # 打开子执行器
            self.left_child.open()
            self.right_child.open()
            
            # 加载并排序左右表数据
            self._load_and_sort_tables()
            
            # 创建合并迭代器
            self.current_iterator = self._create_merge_iterator()
    
    def _load_and_sort_tables(self):
        """加载并排序左右表数据"""
        # 加载左表数据
        while True:
            left_batch = self.left_child.next_batch()
            if not left_batch:
                break
            self.left_sorted_records.extend(left_batch)
        
        # 加载右表数据
        while True:
            right_batch = self.right_child.next_batch()
            if not right_batch:
                break
            self.right_sorted_records.extend(right_batch)
        
        # 排序左表
        self.left_sorted_records.sort(
            key=lambda r: self._extract_join_key(r, self.left_keys)
        )
        
        # 排序右表
        self.right_sorted_records.sort(
            key=lambda r: self._extract_join_key(r, self.right_keys)
        )
        
        self.update_stats(
            rows_processed=len(self.left_sorted_records) + len(self.right_sorted_records)
        )
    
    def _create_merge_iterator(self) -> Iterator[Record]:
        """创建合并迭代器"""
        left_len = len(self.left_sorted_records)
        right_len = len(self.right_sorted_records)
        
        while self.left_index < left_len and self.right_index < right_len:
            self.check_timeout()
            
            left_record = self.left_sorted_records[self.left_index]
            right_record = self.right_sorted_records[self.right_index]
            
            left_key = self._extract_join_key(left_record, self.left_keys)
            right_key = self._extract_join_key(right_record, self.right_keys)
            
            if left_key < right_key:
                # 左表键值较小
                if self.join_type in [JoinType.LEFT, JoinType.FULL]:
                    null_right = self._create_null_record(right_record, "right")
                    merged = self._merge_records(left_record, {})
                    merged.update(null_right)
                    yield merged
                
                self.left_index += 1
            
            elif left_key > right_key:
                # 右表键值较小
                if self.join_type in [JoinType.RIGHT, JoinType.FULL]:
                    # 创建左表的空记录
                    left_template = {}
                    if self.left_sorted_records:
                        # 使用左表第一条记录作为模板创建空记录
                        for key in self.left_sorted_records[0].keys():
                            left_template[key] = None
                    
                    merged = self._merge_records(left_template, right_record)
                    yield merged
                
                self.right_index += 1
            
            else:
                # 键值相等，产生笛卡尔积
                left_start = self.left_index
                right_start = self.right_index
                
                # 找到左表中所有相同键值的记录
                while (self.left_index < left_len and 
                       self._extract_join_key(self.left_sorted_records[self.left_index], self.left_keys) == left_key):
                    
                    # 重置右表索引
                    temp_right_index = right_start
                    
                    # 找到右表中所有相同键值的记录
                    while (temp_right_index < right_len and
                           self._extract_join_key(self.right_sorted_records[temp_right_index], self.right_keys) == right_key):
                        
                        merged = self._merge_records(
                            self.left_sorted_records[self.left_index],
                            self.right_sorted_records[temp_right_index]
                        )
                        self.update_stats(rows_returned=1)
                        yield merged
                        
                        temp_right_index += 1
                    
                    self.left_index += 1
                
                # 更新右表索引
                while (self.right_index < right_len and
                       self._extract_join_key(self.right_sorted_records[self.right_index], self.right_keys) == right_key):
                    self.right_index += 1
        
        # 处理剩余的左表记录（左外连接或全外连接）
        if self.join_type in [JoinType.LEFT, JoinType.FULL]:
            while self.left_index < left_len:
                left_record = self.left_sorted_records[self.left_index]
                null_right = self._create_null_record(
                    self.right_sorted_records[0] if self.right_sorted_records else {}, "right"
                )
                merged = self._merge_records(left_record, {})
                merged.update(null_right)
                yield merged
                self.left_index += 1
        
        # 处理剩余的右表记录（右外连接或全外连接）
        if self.join_type in [JoinType.RIGHT, JoinType.FULL]:
            while self.right_index < right_len:
                right_record = self.right_sorted_records[self.right_index]
                
                # 创建左表的空记录
                left_template = {}
                if self.left_sorted_records:
                    # 使用左表第一条记录作为模板创建空记录
                    for key in self.left_sorted_records[0].keys():
                        left_template[key] = None
                
                merged = self._merge_records(left_template, right_record)
                yield merged
                self.right_index += 1
    
    def close(self):
        """关闭排序合并连接操作符"""
        with self.lock:
            if not self.is_open:
                return
            
            # 关闭子执行器
            self.left_child.close()
            self.right_child.close()
            
            # 清理数据
            self.left_sorted_records.clear()
            self.right_sorted_records.clear()
            
            self.is_open = False
            self.current_iterator = None

